feat(bigframes): Support groupby.agg/transform with udf transpiler#17613
feat(bigframes): Support groupby.agg/transform with udf transpiler#17613TrevorBergeron wants to merge 8 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces subscripting support (via GetItemOp and DynamicGetItemOp) across Ibis, Polars, and SQLGlot compilers for arrays, structs, and strings. It also enables transpilation of custom Python functions (UDFs) within GroupBy aggregations and transformations. The review feedback highlights several critical runtime issues, including missing imports (bigframes and window_specs) that will cause NameErrors, and unhandled string-like inputs in the Polars implementation of GetItemOp.
|
It would've been better if we have broken down this PR into two: one for getItem and the other for dynamicGetItem. But since we are already here, I will accept what it is :) |
| def series_method_to_op(method_name: str) -> Optional[bigframes.operations.RowOp]: | ||
| return SERIES_METHOD_TO_OP.get(method_name) |
There was a problem hiding this comment.
This function might be too trivial. Could we just let the caller perform key lookup directly on the dictionary above?
| def test_series_map_with_struct_and_array_subscript(session): | ||
| import pyarrow as pa | ||
|
|
||
| # Struct setup | ||
| struct_pa_type = pa.struct([("str_field", pa.string()), ("int_field", pa.int64())]) | ||
| pd_struct_series = pd.Series( | ||
| pa.array([{"str_field": "hello", "int_field": 1}], struct_pa_type), | ||
| dtype=pd.ArrowDtype(struct_pa_type), | ||
| ) | ||
| bf_struct_series = bpd.Series(pd_struct_series, session=session) | ||
|
|
||
| # Array setup | ||
| array_pa_type = pa.list_(pa.int64()) | ||
| pd_array_series = pd.Series( | ||
| pa.array([[10, 20]], array_pa_type), | ||
| dtype=pd.ArrowDtype(array_pa_type), | ||
| ) | ||
| bf_array_series = bpd.Series(pd_array_series, session=session) | ||
|
|
||
| # Struct subscripting in UDF | ||
| def get_struct_val(x): | ||
| return x["str_field"] | ||
|
|
||
| bf_struct_res = bf_struct_series.map(get_struct_val).to_pandas() | ||
| pd_struct_res: pd.Series = pd_struct_series.map(get_struct_val) | ||
| assert_series_equal(bf_struct_res, pd_struct_res, check_dtype=False) | ||
|
|
||
| # Array subscripting in UDF | ||
| def get_array_val(x): | ||
| return x[1] | ||
|
|
||
| bf_array_res = bf_array_series.map(get_array_val).to_pandas() | ||
| pd_array_res: pd.Series = pd_array_series.map(get_array_val) | ||
| assert_series_equal(bf_array_res, pd_array_res, check_dtype=False) | ||
|
|
||
| # String setup | ||
| pd_string_series = pd.Series(["hello", "world"]) | ||
| bf_string_series = bpd.Series(pd_string_series, session=session) | ||
|
|
||
| # String subscripting in UDF | ||
| def get_string_val(x): | ||
| return x[1] | ||
|
|
||
| bf_string_res = bf_string_series.map(get_string_val).to_pandas() | ||
| pd_string_res = pd_string_series.map(get_string_val) | ||
| assert_series_equal(bf_string_res, pd_string_res, check_dtype=False) |
There was a problem hiding this comment.
Let's break this test down to multiple smaller tests, each focusing on one subscription type
TrevorBergeron
force-pushed
the
tbergeron_bytecode_groups
branch
from
30aa270 to
6a7eec0
Compare
sycai
left a comment
sycai
left a comment
There was a problem hiding this comment.
It looks like we have many imports at the function block level, which degrades readability. This seems to be a problem from the Jetski side. Could you double check and eliminate those incorrect imports?
| raise NotImplementedError("DataFrameGroupBy.apply is not implemented.") | ||
|
|
||
| def transform(self, func, *args, **kwargs) -> df.DataFrame: | ||
| import bigframes.core.block_transforms as block_transforms |
There was a problem hiding this comment.
It is possible to put this import at the top of the file?
| ) | ||
|
|
||
| def _agg_func(self, func) -> df.DataFrame: | ||
| import bigframes.core.block_transforms as block_transforms |
There was a problem hiding this comment.
put the import statement on top of the file?
|
|
||
| def _agg_dict(self, func: typing.Mapping) -> df.DataFrame: | ||
| aggregations: typing.List[agg_expressions.Aggregation] = [] | ||
| import bigframes.core.block_transforms as block_transforms |
There was a problem hiding this comment.
put statement on top of the file?
| return dataframe if self._as_index else self._convert_index(dataframe) | ||
|
|
||
| def _agg_list(self, func: typing.Sequence) -> df.DataFrame: | ||
| import bigframes.core.block_transforms as block_transforms |
| raise NotImplementedError("SeriesGroupBy.apply is not implemented.") | ||
|
|
||
| def transform(self, func, *args, **kwargs) -> series.Series: | ||
| import bigframes.core.block_transforms as block_transforms |
| ) | ||
|
|
||
| # Scalar context (struct/array getitem ops) | ||
| import bigframes.operations.generic_ops as generic_ops |
| import bigframes.operations.aggregations as agg_ops | ||
|
|
||
| agg_op, _ = agg_ops.lookup_agg_func(attr) | ||
| import bigframes.core.agg_expressions as agg_exprs |
There was a problem hiding this comment.
Please double check if we can move these imports on top of the file
| DATA_DIR / "scalars.jsonl", | ||
| lines=True, | ||
| ) | ||
| from bigframes.testing.utils import convert_pandas_dtypes |
| @pytest.fixture(scope="module", autouse=True) | ||
| def session(): | ||
| import bigframes.core.global_session | ||
| from bigframes.testing import polars_session |
|
|
||
|
|
||
| def test_dataframe_apply_axis_1_with_dynamic_array_subscript(session): | ||
| import pyarrow as pa |
2 participants
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
Fixes #<issue_number_goes_here> 🦕